package org.example.springboot.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.example.springboot.entity.ExpertChat;
import org.example.springboot.repository.ExpertChatRepository;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.stereotype.Component;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Date;
import java.util.List;
import java.util.ArrayList;

/**
 * @author reisen7
 * @date 2025/5/30 13:05
 * @description 
 */

@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {
    // 存储所有连接，key可用userId或expertId
    private static final ConcurrentHashMap<String, List<WebSocketSession>> sessions = new ConcurrentHashMap<>();
    private static ExpertChatRepository expertChatRepository;
    private static final ObjectMapper objectMapper = new ObjectMapper();

    @Autowired
    public void setExpertChatRepository(ExpertChatRepository repo) {
        expertChatRepository = repo;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 解析身份参数
        String userId = getParam(session, "userId");
        String expertId = getParam(session, "expertId");
        String role = getParam(session, "role");
        String key = role.equals("expert") ? "expert_" + expertId : "user_" + userId;
        sessions.computeIfAbsent(key, k -> new ArrayList<>()).add(session);
        System.out.println("[WebSocket] Connected: role=" + role + ", userId=" + userId + ", expertId=" + expertId + ", key=" + key + ", total=" + sessions.get(key).size());
        System.out.println("[WebSocket] All keys: " + sessions.keySet());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 消息格式：{"expertId":1,"userId":2,"content":"hi","fromExpert":true,"sendTime":"..."}
        String payload = message.getPayload();
        ExpertChat chat = objectMapper.readValue(payload, ExpertChat.class);
        chat.setSendTime(new Date());
        // 保存到 expert_chat 表
        expertChatRepository.save(chat);
        // 转发给目标
        String targetKey = chat.getFromExpert() ? "user_" + chat.getUserId() : "expert_" + chat.getExpertId();
        String json = objectMapper.writeValueAsString(chat);
        List<WebSocketSession> targetSessions = sessions.get(targetKey);
        System.out.println("[WebSocket] Message: fromExpert=" + chat.getFromExpert() + ", userId=" + chat.getUserId() + ", expertId=" + chat.getExpertId() + ", targetKey=" + targetKey + ", sessions.keys=" + sessions.keySet());
        if (targetSessions != null) {
            for (WebSocketSession ws : targetSessions) {
                if (ws.isOpen()) {
                    ws.sendMessage(new TextMessage(json));
                }
            }
            System.out.println("[WebSocket] Push to " + targetKey + ", count: " + targetSessions.size());
        } else {
            System.out.println("[WebSocket] No session for targetKey: " + targetKey);
        }
        // 回显给自己
        if (session.isOpen()) {
            session.sendMessage(new TextMessage(json));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        // 移除
        sessions.forEach((key, list) -> {
            if (list.remove(session)) {
                System.out.println("[WebSocket] Disconnected: " + key + ", remain: " + list.size());
            }
        });
        System.out.println("[WebSocket] All keys after disconnect: " + sessions.keySet());
    }

    private String getParam(WebSocketSession session, String key) {
        String query = session.getUri().getQuery();
        for (String param : query.split("&")) {
            String[] kv = param.split("=");
            if (kv.length == 2 && kv[0].equals(key)) return kv[1];
        }
        return null;
    }
} 